import math
import random
import matplotlib.pyplot as plt
from typing import List, Tuple, Union

# 设置中文字体为黑体，用来正常显示中文标签，并设置能正常显示负号
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# 加载Iris数据集
def load_iris_data() -> List[Tuple[List[float], str]]:
    iris_data = []
    try:
        with open('iris.data', 'r') as file_obj:
            for line in file_obj:
                if line.strip():
                    row = line.strip().split(',')
                    feature_values = list(map(float, row[:-1]))
                    label_value = row[-1]
                    iris_data.append((feature_values, label_value))
    except FileNotFoundError:
        print("数据集文件 'iris.data' 不存在，请检查文件路径！")
    return iris_data


# 计算欧氏距离
def calculate_euclidean_distance(point_1: List[float], point_2: List[float]) -> float:
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(point_1, point_2)))


# KNN分类算法
def knn_classification(train_dataset: List[Tuple[List[float], str]], test_point: List[float], neighbor_count: int) -> str:
    distance_list = []
    for features, label in train_dataset:
        distance = calculate_euclidean_distance(features, test_point)
        distance_list.append((distance, label))
    distance_list.sort(key=lambda x: x[0])  # 按距离排序
    nearest_neighbors = distance_list[:neighbor_count]
    label_list = [label for _, label in nearest_neighbors]
    return max(set(label_list), key=label_list.count)  # 多数表决


# 划分训练集和测试集
def split_train_test_data(data: List[Tuple[List[float], str]], test_proportion: float = 0.2) -> Tuple[
    List[Tuple[List[float], str]], List[Tuple[List[float], str]]]:
    random.shuffle(data)
    split_index = int(len(data) * (1 - test_proportion))
    train_set = data[:split_index]
    test_set = data[split_index:]
    return train_set, test_set


# 可视化分类结果
def visualize_classification_results(test_dataset: List[Tuple[List[float], str]], prediction_results: List[str], neighbor_count: int) -> None:
    color_mapping = {'Iris-setosa': 'r', 'Iris-versicolor': 'g', 'Iris-virginica': 'b'}
    marker_mapping = {'correct': 'o', 'wrong': 's'}
    plt.figure(figsize=(10, 6))
    for (features, true_label), pred_label in zip(test_dataset, prediction_results):
        is_correct = true_label == pred_label
        plt.scatter(features[0], features[1], c=color_mapping[true_label],
                    marker=marker_mapping['correct' if is_correct else 'wrong'])
    plt.title(f'KNN分类结果 (k={neighbor_count})')
    plt.xlabel('特征 1')
    plt.ylabel('特征 2')
    plt.legend(loc='upper right', labels=['山鸢尾', '变色鸢尾', '维吉尼亚鸢尾'])
    plt.show()


# 主函数
if __name__ == "__main__":
    # 加载数据
    iris_data = load_iris_data()
    if not iris_data:
        exit(1)
    train_set, test_set = split_train_test_data(iris_data)

    neighbor_count = 3
    prediction_results = []
    correct_count = 0
    total_count = len(test_set)
    class_correct_counts_dict = {'Iris-setosa': 0, 'Iris-versicolor': 0, 'Iris-virginica': 0}
    class_total_counts_dict = {'Iris-setosa': 0, 'Iris-versicolor': 0, 'Iris-virginica': 0}

    chinese_label_mapping = {
        'Iris-setosa': '山鸢尾',
        'Iris-versicolor': '变色鸢尾',
        'Iris-virginica': '维吉尼亚鸢尾'
    }

    for features, label in test_set:
        prediction = knn_classification(train_set, features, neighbor_count)
        prediction_results.append(prediction)
        if prediction == label:
            correct_count += 1
            class_correct_counts_dict[label] += 1
        class_total_counts_dict[label] += 1

    accuracy = correct_count / total_count
    print(f"KNN分类准确率: {accuracy * 100:.2f}%")
    print("各类别分类情况：")
    for class_name in class_correct_counts_dict:
        precision = class_correct_counts_dict[class_name] / class_total_counts_dict[class_name] if class_total_counts_dict[
                                                                                                 class_name] > 0 else 0
        chinese_label = chinese_label_mapping[class_name]
        print(f"{chinese_label}类别准确率: {precision * 100:.2f}%")

    # 可视化分类结果
    visualize_classification_results(test_set, prediction_results, neighbor_count)